#-*- coding: utf-8 -*-  
  
import numpy as np  
import os, re  
from PIL import Image, ImageEnhance, ImageOps  
from sklearn.model_selection import train_test_split  
from sklearn.neural_network import MLPClassifier  
from sklearn.metrics import confusion_matrix, accuracy_score  
from skimage.feature import hog, graycomatrix, graycoprops  
from skimage.color import rgb2gray  
from skimage.transform import resize  
from skimage.exposure import rescale_intensity  
import skimage.io as io  
  
# 图像切割及特征提取  
path = 'C:/Code/Code_Python/数据挖掘/作业3/代码/data/images/'  # 图片所在路径  
  
# 自定义获取图片名称函数  
def get_img_names(path=path):  
    filenames = os.listdir(path)  
    img_names = [i for i in filenames if re.findall(r'^\d_\d+\.jpg$', i)]  
    return img_names  
  
# 自定义图像增强函数  
def enhance_image(img):  
    enhancer = ImageEnhance.Brightness(img)  
    img = enhancer.enhance(1.2)  # 增加亮度  
    img = img.convert('L')  # 转为灰度图  
    return img  
  
# 额外数据增强（旋转、翻转）  
def augment_image(img):  
    imgs = [img]  
    imgs.append(img.rotate(90, expand=True))  
    imgs.append(img.rotate(180, expand=True))  
    imgs.append(img.rotate(270, expand=True))  
    imgs.append(ImageOps.flip(img))  
    return imgs  
  
# 批量处理图片数据  
img_names = get_img_names(path=path)  
n = len(img_names)  
data = []  
labels = []  
  
for i in range(n):  
    img = Image.open(path + img_names[i])  
    M, N = img.size  
    img = img.crop((M//2-50, N//2-50, M//2+50, N//2+50))  # 图片切割  
    img = enhance_image(img)  # 图像增强  
  
    augmented_imgs = augment_image(img)  
    for aug_img in augmented_imgs:  
        # 将图像转换为灰度图并转换为numpy数组  
        gray_img = np.array(aug_img, dtype=np.uint8)  # 直接以uint8读取，避免之后的转换  
  
        # 使用 skimage 计算 GLCM 和纹理特征  
        glcm = graycomatrix(gray_img, [1, 2], [0, np.pi/4, np.pi/2, 3*np.pi/4], 256, symmetric=True, normed=True)  
        contrast = graycoprops(glcm, 'contrast').mean()  
        dissimilarity = graycoprops(glcm, 'dissimilarity').mean()  
        asm = graycoprops(glcm, 'ASM').mean()  
        energy = graycoprops(glcm, 'energy').mean()  
  
        # 计算 HOG 特征  
        hog_features, hog_image = hog(gray_img, pixels_per_cell=(16, 16),  
                                      cells_per_block=(2, 2), visualize=True)  
  
        # 将提取的特征存储到数据数组中  
        feature_vector = np.array([contrast, dissimilarity, asm, energy] + hog_features.flatten().tolist())  
        data.append(feature_vector)  
        labels.append(img_names[i][0])  # 样本标签  
  
data = np.array(data)  
labels = np.array(labels, dtype=np.str_)  
  
# 数据拆分,训练集、测试集  
data_tr, data_te, label_tr, label_te = train_test_split(data, labels, test_size=0.2, random_state=10)  
  
model = MLPClassifier(hidden_layer_sizes=(100,), max_iter=300, random_state=5).fit(data_tr, label_tr)  
  
# 水质评价  
pre_te = model.predict(data_te)  
cm_te = confusion_matrix(label_te, pre_te)  
print(cm_te)  
print(accuracy_score(label_te, pre_te))